"""
之前一直想尝试的一件事情：
对于一个二维的正态分布，采样一些点作为样本，加上噪声，然后使用MLP作为网络去拟合,是否可以拟合出来这样的分布？
简而言之就是拟合二维正态分布的概率密度函数
"""
from re import X
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
from scipy.stats import multivariate_normal
import sys


def show_multivariate_normal():
    x = np.linspace(-1, 3, 100)
    y = np.linspace(0, 4, 100)
    X, Y = np.meshgrid(x, y)
    pos = np.dstack((X, Y))
    print(pos)
    print('pos shape:', pos.shape)
    mu = np.array([1, 2])
    cov = np.array([[.5, .25],[.25, .5]])
    rv = multivariate_normal(mu, cov) 
    Z = rv.pdf(pos) # pdf 概率密度函数 probability density function
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    X, Y, Z = X*10, Y*10, Z*10
    print('X shape:{}, Y shape:{}, Z shape:{}'.format(X.shape, Y.shape, Z.shape))
    ax.plot_surface(X, Y, Z)
    #fig.show() # fig.show()会一闪而过
    plt.show()

def genData():
    mu = np.array([1, 2]) # 二维正态分布的mean
    cov = np.array([[.5, .25],[.25, .5]]) # 二维正态分布的协方差矩阵
    rv = multivariate_normal(mu, cov) # 获取二维正态分布

    n = 2e4

    # x y 的数据均匀分布在 [-2, 4] in x / [-1, 5] in y
    x = -2. + 6. * np.random.rand(int(n))
    #print(X[:5])
    y = -1. + 6. * np.random.rand(int(n))

    noise = np.random.randn(int(n)) - 0.5
    noise = noise * 0.1

    pos = np.dstack((x, y))
    z = rv.pdf(pos)
    print(z.shape)

    z = z * 10
    # z += noise
    xyz = np.dstack((x, y, z))
    xyz = np.squeeze(xyz)
    print(xyz.shape)
    
    plt.figure("3D Scatter", facecolor="lightgray")
    ax3d = plt.gca(projection="3d")  # 创建三维坐标

    plt.title('3D Scatter', fontsize=20)
    ax3d.set_xlabel('x', fontsize=14)
    ax3d.set_ylabel('y', fontsize=14)
    ax3d.set_zlabel('z', fontsize=14)
    plt.tick_params(labelsize=10)

    ax3d.scatter(x, y, z,  cmap="jet", marker="o")

    plt.show()


    # np.save('./multivariate_normal_data', xyz)


def show_scatter():
    path = r'Z:\FitNormal\Data\Result_MSELoss.npy'
    data = np.load(path)
    x, y, z = data[:, 0], data[:, 1], data[:, 2]

    plt.figure("3D Scatter", facecolor="lightgray")
    ax3d = plt.gca(projection="3d")  # 创建三维坐标

    plt.title('3D Scatter', fontsize=20)
    ax3d.set_xlabel('x', fontsize=14)
    ax3d.set_ylabel('y', fontsize=14)
    ax3d.set_zlabel('z', fontsize=14)
    plt.tick_params(labelsize=10)

    ax3d.scatter(x, y, z,  cmap="jet", marker="o")

    plt.show()

    

if __name__ == '__main__':

    # show_multivariate_normal()
    # sys.exit()
    # genData()
    show_scatter()

    